# coding=utf-8
##
## Author: jmdvirus@aliyun.com
##
## Create: 2019年02月14日 星期四 15时39分35秒
##

import numpy as np
import os
import gzip
import struct
import logging
import mxnet as mx
import matplotlib.pyplot as plt

logging.getLogger().setLevel(logging.DEBUG)

def read_data(label_url, image_url):
    with gzip.open(label_url) as flbl:
        magic, num = struct.unpack(">II", flbl.read(8))
        label = np.fromstring(flbl.read(), dtype=np.int8)
    with gzip.open(image_url, 'rb') as fimg:
        magic, num, rows, cols = struct.unpack(">IIII", fimg.read(16))
        image = np.fromstring(fimg.read(), dtype=np.uint8)
        image = image.reshape(len(label), 1, rows, cols)
        image = image.astype(np.float32)/255.0

    return (label, image)

prefix_path = "/opt/data/proj/src/ai/fashion-mnist/data/fashion/"

(train_lbl, train_img) = read_data(prefix_path + 'train-labels-idx1-ubyte.gz', 
        prefix_path + 'train-images-idx3-ubyte.gz')
(val_lbl, val_img) = read_data(prefix_path + 't10k-labels-idx1-ubyte.gz', 
        prefix_path + 't10k-images-idx3-ubyte.gz')

batch_size = 32

def cba(src, suffix, num_filter, kernel, pad):
    conv = mx.sym.Convolution(data=src, name="conv"+suffix, kernel=(kernel, kernel), 
            pad=(pad,pad), num_filter=num_filter)
    bn = mx.sym.BatchNorm(data=conv, name="bn"+suffix, fix_gamma=False)
    act = mx.sym.Activation(data=bn, name="act"+suffix, act_type="relu")
    return act

def layer(src, layer, num_filter, pad):
    conv1 = cba(src, layer+"1", num_filter, 3, pad)
    conv2 = cba(conv1, layer+"2", num_filter, 3, pad)
    pool = mx.sym.Pooling(data=conv2, name="pool"+layer, pool_type="max",
            kernel=(2,2), stride=(2,2))
    return pool

data = mx.symbol.Variable('data')
net = layer(data, "1", 32, 1)
net = layer(net, "2", 64, 1)
net = layer(net, "3", 64, 1)
net = cba(net, "4", 128, 3, 0)

net = mx.sym.Convolution(data=net, name="final", kernel=(1,1), num_filter=10)

net = mx.sym.Flatten(data=net, name="flatten")
net = mx.sym.SoftmaxOutput(data=net, name='softmax')

shape = {"data": (batch_size, 1, 28, 28)}
mx.viz.print_summary(symbol=net, shape=shape)

def shownet(net):
    mx.viz.plot_network(symbol=net, shape=shape).view()

def totrain(net):
    module = mx.mod.Module(symbol=net, context=mx.cpu(0))

    val_iter = mx.io.NDArrayIter(val_img, val_lbl, batch_size)

    for epoch in range(2):
        aug_img = train_img.copy()
        for i in range(aug_img.shape[0]):
            if np.random.random() < 0.5:
                aug_img[i][0] = np.fliplr(aug_img[i][0])
            amt = np.random.randint(0,3)
            if amt > 0:
                if np.random.random() < 0.5:
                    aug_img[i][0] = np.pad(aug_img[i][0], ((0,0), (amt, 0)),
                        mode='constant')[:, :-amt]
                else:
                    aug_img[i][0] = np.pad(aug_img[i][0], ((0,0), (0, amt)),
                            mode='constant')[:, amt:]

            amt = np.random.randint(0, 3)
            if amt > 0:
                if np.random.random() < 0.5:
                    aug_img[i][0] = np.pad(aug_img[i][0], ((amt, 0), (0, 0)), 
                            mode='constant')[:-amt, :]
                else:
                    aug_img[i][0] = np.pad(aug_img[i][0], ((0, amt), (0, 0)),
                            mode='constant')[amt:, :]

            x_size = np.random.randint(1, 6)
            y_size = np.random.randint(1, 6)
            x_begin = np.random.randint(0, 28-x_size + 1)
            y_begin = np.random.randint(0, 28-y_size + 1)
            aug_img[i][0][x_begin:x_begin + x_size, y_begin:y_begin+y_size] = 0

        train_iter = mx.io.NDArrayIter(aug_img, train_lbl, batch_size, shuffle=True)

        lr = 0.06 * pow(0.95, epoch)

        print("epoch " + str(epoch) + ", learning_rate = " + str(lr))

        module.fit(
                train_iter,
                eval_data = val_iter, 
                optimizer = 'sgd',
                optimizer_params = {'learning_rate': lr},
                num_epoch = 1,
                batch_end_callback = mx.callback.Speedometer(batch_size, 60000/batch_size)
                )

if __name__ == "__main__":
    shownet(net)
    #totrain(net)

